import os
import sys
import argparse
import cv2
import numpy as np
import yaml
from collections import defaultdict
from discret_image import active_size_map_discretization


MODES = [
    "pgm_to_png",
    "png_red_to_pgm",
    "node_edge_mode",
    "auto_mode",
    "yml2txt",
    "remove_blocked_edges",
    "png_green_to_pgm",
]


def parse_arguments():
    """Parse command-line arguments with flexible arity support."""
    parser = argparse.ArgumentParser(
        description="Map processing script with flexible parameters."
    )
    parser.add_argument(
        "root",
        nargs="?",
        default="../datas/",
        help="Root directory (default: ../datas/)",
    )
    parser.add_argument(
        "yaml_file_name",
        nargs="?",
        default="outdoor_final.yaml",
        help="YAML file name (default: outdoor_final.yaml)",
    )
    parser.add_argument(
        "node_radius", nargs="?", type=int, default=50, help="Node radius (default: 50)"
    )
    return parser.parse_args()


def initialize_paths(root, yaml_file_name, node_radius):
    """Initialize all file paths based on parameters."""
    map_file_name = yaml_file_name.replace(".yaml", ".pgm")
    pgm_path = os.path.join(root, map_file_name)
    config_path = os.path.join(root, yaml_file_name)
    png_path = pgm_path.replace(".pgm", ".png")
    graph_path = pgm_path.replace(".pgm", "_graph.yaml")
    vis_png_path = pgm_path.replace(".pgm", "_vis.png")
    return pgm_path, config_path, png_path, graph_path, vis_png_path, node_radius


def load_config(config_path):
    """Load YAML config with error handling."""
    if not os.path.exists(config_path):
        raise ValueError(f"Config file not found: {config_path}")
    with open(config_path, "r") as f:
        return yaml.safe_load(f)


def load_image(image_path, flags=cv2.IMREAD_UNCHANGED):
    """Load image with error handling."""
    img = cv2.imread(image_path, flags)
    if img is None:
        raise ValueError(f"Unable to load image: {image_path}")
    return img


def coordinate_transform(config, nodes, resolution, origin, height, inverse=False):
    """
    Transform nodes between world and image coordinates.
    inverse=True: image -> world; else: world -> image.
    """
    x, y, yaw = origin
    yaw_rad = np.deg2rad(yaw)
    R = np.array(
        [[np.cos(yaw_rad), -np.sin(yaw_rad)], [np.sin(yaw_rad), np.cos(yaw_rad)]]
    )
    T = np.array([x, y])

    transformed = []
    for node in nodes:
        world_pt = np.array(node[:2])
        if inverse:
            # Image to world (with Y flip)
            img_pt = np.array([node[0], height - node[1]])
            world_pt = R @ (img_pt * resolution) + T
            transformed.append([*world_pt, 0])
        else:
            # World to image (with Y flip)
            img_pt = R.T @ ((world_pt - T) / resolution)
            x_img, y_img = img_pt
            transformed.append([x_img, height - y_img])
    return transformed


def turn_pgm_to_png(pgm_path, png_path):
    """Convert PGM to PNG."""
    img = load_image(pgm_path)
    if len(img.shape) == 2:
        img_rgb = cv2.cvtColor(img, cv2.COLOR_GRAY2RGB)
    else:
        img_rgb = img[:, :, :3].copy()
    cv2.imwrite(png_path, img_rgb)
    print(f"Converted {pgm_path} to {png_path}")


def turn_png_channel_to_pgm(png_path, pgm_path, low_bound, high_bound, invert=False):
    """Generic PNG channel extraction to PGM (red or green)."""
    img = load_image(png_path)
    mask = cv2.inRange(img, low_bound, high_bound)
    masked = cv2.bitwise_and(img, img, mask=mask)
    gray = cv2.cvtColor(masked, cv2.COLOR_BGR2GRAY)
    gray[gray > 0] = 255
    if invert:
        gray = 255 - gray
    cv2.imwrite(pgm_path, gray)
    print(f"Extracted channel from {png_path} to {pgm_path}")


def turn_png_red_to_pgm(png_path, pgm_path):
    """Convert red channel in PNG to PGM."""
    low = np.array([0, 0, 200])
    high = np.array([50, 50, 255])
    turn_png_channel_to_pgm(png_path, pgm_path, low, high)


def turn_png_green_to_pgm(png_path, pgm_path):
    """Convert green channel in PNG to PGM (inverted for obstacles)."""
    low = np.array([0, 200, 0])
    high = np.array([20, 255, 20])
    turn_png_channel_to_pgm(png_path, pgm_path, low, high, invert=True)


def draw_graph_on_image(img, graph):
    """Draw graph on image and display."""
    nodes = np.array(graph["nodes"], dtype=np.int32)
    edges = graph["edges"]

    if len(img.shape) == 2:
        img_bgr = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)
    else:
        img_bgr = img.copy()

    # Draw nodes
    for idx in range(len(nodes)):
        cv2.circle(img_bgr, tuple(nodes[idx][:2]), 1, (0, 255, 0), -1)

    # Draw edges (vectorized where possible, but loop for simplicity)
    for e in edges:
        pt1 = tuple(nodes[e[0]][:2])
        pt2 = tuple(nodes[e[1]][:2])
        cv2.line(img_bgr, pt1, pt2, (255, 0, 255), 1)

    cv2.imshow("Graph", img_bgr)
    cv2.waitKey(0)
    cv2.destroyAllWindows()
    return img_bgr


def get_graph_from_image(png_path, config_path, graph_path):
    """Extract graph from image with red nodes and blue edges."""
    img = load_image(png_path)
    config = load_config(config_path)
    resolution = config["resolution"]
    origin = config["origin"]
    height, width = img.shape[:2]

    # Node detection (red dots)
    lower_red = np.array([0, 0, 200])
    upper_red = np.array([0, 0, 255])
    mask_red = cv2.inRange(img, lower_red, upper_red)
    contours, _ = cv2.findContours(mask_red, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    nodes = []
    for cnt in contours:
        M = cv2.moments(cnt)
        if M["m00"] != 0:
            cx = int(M["m10"] / M["m00"])
            cy = int(M["m01"] / M["m00"])
            nodes.append([cx, cy])

    # Edge detection (blue lines) - optimized with connected components
    lower_blue = np.array([200, 0, 0])
    upper_blue = np.array([255, 0, 0])
    mask_blue = cv2.inRange(img, lower_blue, upper_blue)
    kernel = np.ones((3, 3), np.uint8)
    mask_blue = cv2.morphologyEx(mask_blue, cv2.MORPH_CLOSE, kernel)

    num_labels, labels = cv2.connectedComponents(mask_blue)
    label_to_points = defaultdict(list)
    points = cv2.findNonZero(mask_blue)
    if points is not None:
        points = points.reshape(-1, 2)
        for pt in points:
            x, y = int(pt[0]), int(pt[1])
            label = labels[y, x]
            label_to_points[label].append([x, y])

    edges = []
    nodes_np = np.array(nodes)
    for label in range(1, num_labels):
        cluster = np.array(label_to_points[label], dtype=np.float32)
        if len(cluster) < 2:
            continue

        # PCA for line direction
        mean, eigenvectors, _ = cv2.PCACompute(cluster, mean=np.array([]))
        direction = eigenvectors[0]

        # Project and find endpoints
        vectors = cluster - mean
        projections = np.dot(vectors, direction)
        min_idx, max_idx = np.argmin(projections), np.argmax(projections)
        pt1, pt2 = cluster[min_idx].astype(int), cluster[max_idx].astype(int)

        # Find nearest nodes (vectorized distance check)
        dists1 = np.linalg.norm(nodes_np - pt1, axis=1)
        dists2 = np.linalg.norm(nodes_np - pt2, axis=1)
        nearest1 = np.argmin(dists1)
        nearest2 = np.argmin(dists2)

        if nearest1 != nearest2:  # Avoid self-loops
            edges.append([nearest1, nearest2])

    graph = {"nodes": nodes, "edges": edges}
    draw_graph_on_image(img, graph)

    # Transform and save
    graph["nodes"] = coordinate_transform(
        config, graph["nodes"], resolution, origin, height, inverse=True
    )
    save_graph_to_yaml(graph, config_path, graph_path)
    return graph


def preprocess_image(img):
    """Preprocess grayscale image for discretization."""
    print("[Preprocessing] Denoising and enhancing...")
    denoised = cv2.medianBlur(img, 3)
    kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (3, 3))
    closed = cv2.morphologyEx(denoised, cv2.MORPH_CLOSE, kernel, iterations=1)
    smoothed = cv2.GaussianBlur(closed, (3, 3), 0)
    _, binary = cv2.threshold(smoothed, 150, 255, cv2.THRESH_BINARY)
    return binary


def flood_fill_from_root(binary_img, root_pixel):
    """Flood fill from root pixel to mark walkable areas."""
    h, w = binary_img.shape
    mask = np.zeros((h + 2, w + 2), dtype=np.uint8)
    flooded = binary_img.copy()
    flood_color = 100
    lo_diff, hi_diff = 10, 10
    flags = 4 | (lo_diff << 8) | cv2.FLOODFILL_FIXED_RANGE
    cv2.floodFill(flooded, mask, root_pixel, flood_color, lo_diff, hi_diff, flags)
    return flooded, flood_color


def remove_isolated_nodes(graph):
    """Remove isolated nodes and update edges."""
    print("[Optimization] Removing isolated nodes...")
    nodes = graph["nodes"]
    edges = graph["edges"]
    degrees = np.zeros(len(nodes), dtype=int)
    for u, v in edges:
        degrees[u] += 1
        degrees[v] += 1

    isolated = np.where(degrees == 0)[0]
    print(f"  >> Isolated nodes: {len(isolated)}")

    keep_mask = ~np.isin(range(len(nodes)), isolated)
    new_nodes = [nodes[i] for i in range(len(nodes)) if keep_mask[i]]
    old_to_new = {old: new for new, old in enumerate(np.where(keep_mask)[0])}

    new_edges = [
        [old_to_new[u], old_to_new[v]]
        for u, v in edges
        if keep_mask[u] and keep_mask[v]
    ]

    graph["nodes"] = new_nodes
    graph["edges"] = new_edges
    print(f"  >> Updated: {len(new_nodes)} nodes, {len(new_edges)} edges")
    return graph


def get_graph_by_discretization(
    img_path, config_path, graph_path, vis_png_path, node_radius, root=[0, 0]
):
    """Generate graph via image discretization."""
    img = load_image(img_path, cv2.IMREAD_GRAYSCALE)
    config = load_config(config_path)
    resolution = config["resolution"]
    origin = config["origin"]
    height, width = img.shape[:2]

    # Compute root pixel
    yaw = origin[2]
    R = np.array(
        [[np.cos(yaw), -np.sin(yaw)], [np.sin(yaw), np.cos(yaw)]], dtype=np.float32
    )
    T = np.array([origin[0], origin[1]], dtype=np.float32)
    root_world = np.array(root[:2], dtype=np.float32)
    map_xy = (R.T @ (root_world - T)) / resolution
    root_pixel = np.round([map_xy[0], height - 1 - map_xy[1]]).astype(int)
    root_pixel = np.clip(root_pixel, [0, 0], [width - 1, height - 1]).tolist()

    binary_img = preprocess_image(img)
    flooded, runnable_color = flood_fill_from_root(binary_img, tuple(root_pixel))

    result_img, graph = active_size_map_discretization(
        flooded, runnable_color=runnable_color, node_radius=node_radius
    )

    graph = remove_isolated_nodes(graph)

    cv2.imwrite(vis_png_path, result_img)
    save_graph_to_yaml(graph, config_path, graph_path)
    return graph


def save_graph_to_yaml(graph, config_path, graph_path):
    """Save graph to YAML after coordinate transform."""

    # 转换所有numpy数据类型为原生Python类型
    def convert_numpy(obj):
        if isinstance(obj, np.ndarray):
            return obj.tolist()
        elif isinstance(obj, np.integer):
            return int(obj)
        elif isinstance(obj, np.floating):
            return float(obj)
        elif isinstance(obj, list):
            return [convert_numpy(item) for item in obj]
        return obj

    # 在保存前转换数据
    converted_graph = {
        "nodes": convert_numpy(graph["nodes"]),
        "edges": convert_numpy(graph["edges"]),
    }

    config = load_config(config_path)
    resolution = config["resolution"]
    origin = config["origin"]
    height = load_image(
        os.path.join(os.path.dirname(config_path), config_path.replace(".yaml", ".png"))
    ).shape[0]

    converted_graph["nodes"] = convert_numpy(
        coordinate_transform(
            config, converted_graph["nodes"], resolution, origin, height, inverse=True
        )
    )

    with open(graph_path, "w") as f:
        yaml.dump(converted_graph, f, default_flow_style=False)
    print(f"Saved graph to {graph_path}")


def visualize_graph(img_path, graph_path, config_path):
    """Visualize graph on image."""
    img = load_image(img_path)
    with open(graph_path, "r") as f:
        graph = yaml.safe_load(f)
    config = load_config(config_path)
    resolution = config["resolution"]
    origin = config["origin"]
    height = img.shape[0]
    graph["nodes"] = coordinate_transform(
        config, graph["nodes"], resolution, origin, height
    )
    draw_graph_on_image(img, graph)


def yml2txt(graph_path, save_dir):
    """Convert YAML graph to TXT files (points and adj matrix)."""
    with open(graph_path, "r") as f:
        graph = yaml.safe_load(f)
    num_nodes = len(graph["nodes"])

    # Save points
    point_path = os.path.join(save_dir, "points.txt")
    with open(point_path, "w") as f:
        for node in graph["nodes"]:
            f.write(f"{node[0]} {node[1]}\n")

    # Adjacency matrix
    adj = np.eye(num_nodes, dtype=int)
    for u, v in graph["edges"]:
        adj[u, v] = adj[v, u] = 1

    adj_path = os.path.join(save_dir, "adj_matrix.txt")
    with open(adj_path, "w") as f:
        f.write(f"{num_nodes}\n")
        np.savetxt(f, adj, fmt="%d", delimiter=" ")

    print(f"Saved TXT files to {save_dir}")


def remove_blocked_nodes_edges(png_path, config_path, graph_path):
    """Remove nodes/edges in blocked areas."""
    img = load_image(png_path)
    config = load_config(config_path)
    resolution = config["resolution"]
    origin = config["origin"]
    height, width = img.shape[:2]

    with open(graph_path, "r") as f:
        graph = yaml.safe_load(f)

    # Transform to image coords
    img_nodes = coordinate_transform(config, graph["nodes"], resolution, origin, height)
    valid_nodes = []
    valid_indices = []

    for idx, (x_img, y_img) in enumerate(img_nodes):
        x_int, y_int = int(round(x_img)), int(round(y_img))
        if 0 <= x_int < width and 0 <= y_int < height:
            pixel = img[y_int, x_int]
            is_valid = (
                pixel > 0 if len(img.shape) == 2 else not (pixel == [0, 0, 0]).all()
            )
            if is_valid:
                valid_nodes.append(graph["nodes"][idx])
                valid_indices.append(idx)

    valid_set = set(valid_indices)
    valid_edges = []
    for u, v in graph["edges"]:
        if u in valid_set and v in valid_set:
            new_u = valid_indices.index(u)
            new_v = valid_indices.index(v)
            valid_edges.append([new_u, new_v])

    new_graph = {"nodes": valid_nodes, "edges": valid_edges}
    with open(graph_path, "w") as f:
        yaml.dump(new_graph, f, default_flow_style=False)

    print(
        f"Cleaned: {len(graph['nodes'])} -> {len(valid_nodes)} nodes; {len(graph['edges'])} -> {len(valid_edges)} edges"
    )


def run_mode(
    mode,
    pgm_path,
    png_path,
    config_path,
    graph_path,
    vis_png_path,
    node_radius,
    root=[0, 0],
):
    """Run a single mode."""
    print(f"Running mode: {mode}")
    if mode == "pgm_to_png":
        turn_pgm_to_png(pgm_path, png_path)
    elif mode == "png_red_to_pgm":
        turn_png_red_to_pgm(png_path, pgm_path)
    elif mode == "node_edge_mode":
        get_graph_from_image(png_path, config_path, graph_path)
    elif mode == "auto_mode":
        get_graph_by_discretization(
            png_path, config_path, graph_path, vis_png_path, node_radius, root
        )
    elif mode == "yml2txt":
        yml2txt(graph_path, os.path.dirname(graph_path))
    elif mode == "remove_blocked_edges":
        remove_blocked_nodes_edges(png_path, config_path, graph_path)
    elif mode == "png_green_to_pgm":
        turn_png_green_to_pgm(png_path, pgm_path)
    else:
        visualize_graph(png_path, graph_path, config_path)


def process_new_map(
    pgm_path, png_path, config_path, graph_path, vis_png_path, node_radius
):
    """Process new map workflow."""
    run_mode(
        "pgm_to_png",
        pgm_path,
        png_path,
        config_path,
        graph_path,
        vis_png_path,
        node_radius,
    )
    run_mode(
        "auto_mode",
        pgm_path,
        png_path,
        config_path,
        graph_path,
        vis_png_path,
        node_radius,
    )
    run_mode(
        "yml2txt",
        pgm_path,
        png_path,
        config_path,
        graph_path,
        vis_png_path,
        node_radius,
    )


def process_old_map(png_path, config_path, graph_path):
    """Process old map workflow."""
    run_mode("png_red_to_pgm", None, png_path, config_path, graph_path, None, None)
    run_mode(
        "remove_blocked_edges", None, png_path, config_path, graph_path, None, None
    )


if __name__ == "__main__":
    args = parse_arguments()
    pgm_path, config_path, png_path, graph_path, vis_png_path, node_radius = (
        initialize_paths(args.root, args.yaml_file_name, args.node_radius)
    )

    # Default run as per original
    # run_mode(
    #     "png_red_to_pgm",
    #     None,
    #     png_path,
    #     config_path,
    #     graph_path,
    #     vis_png_path,
    #     node_radius,
    # )
    process_new_map(
        pgm_path, png_path, config_path, graph_path, vis_png_path, node_radius
    )
